import zipfile
import shutil
import os
import tempfile

def cv2_modify_whl_version_py(whl_path, output_whl_path=None):
    if output_whl_path is None:
        output_whl_path = whl_path

    # 创建临时解压目录
    with tempfile.TemporaryDirectory() as tmpdir:
        with zipfile.ZipFile(whl_path, 'r') as zip_ref:
            zip_ref.extractall(tmpdir)

        # 找到 version.py 文件路径
        version_py_path = None
        for root, dirs, files in os.walk(tmpdir):
            for file in files:
                if file == "version.py":
                    version_py_path = os.path.join(root, file)
                    break
            if version_py_path:
                break

        if not version_py_path:
            raise FileNotFoundError("version.py not found in the .whl package")

        # 修改 version.py 文件内容
        with open(version_py_path, "r") as f:
            content = f.read()

        new_content = []
        for line in content.splitlines():
            if line.startswith("ci_build = "):
                new_content.append("ci_build = True")
            else:
                new_content.append(line)

        with open(version_py_path, "w") as f:
            f.write("\n".join(new_content) + "\n")

        # 重新打包成 .whl
        with zipfile.ZipFile(output_whl_path, 'w', zipfile.ZIP_DEFLATED) as new_zip:
            for foldername, subfolders, filenames in os.walk(tmpdir):
                for filename in filenames:
                    filepath = os.path.join(foldername, filename)
                    arcname = os.path.relpath(filepath, tmpdir)
                    new_zip.write(filepath, arcname)

    print(f"✅ Modified wheel saved to: {output_whl_path}")


def fix_some_whl(pkg_full_path):
    if 'opencv_python' in pkg_full_path and 'headless' not in pkg_full_path:
        cv2_modify_whl_version_py(pkg_full_path)


# 示例用法
cv2_modify_whl_version_py("/home/zq/whl_build/opencv-python-4.12.0.88/wheelhouse/opencv_python-4.12.0.88-cp313-cp313-linux_x86_64.whl")
